//
// Created by hamlet on 23-3-12.
//

#include "KNNSolver"

constexpr double SKIP_VALUE = 0.0001;

namespace FansionML {
    KNNSolver::KNNSolver(int k) : m_k(k), m_data(), m_label() {}

    // 根据KNN算法的原理，KNN不需要预训练
    void KNNSolver::train(const DataT &data, const LabelT &label) {
        m_data = &data;
        m_label = &label;
        m_dis.resize(m_data->rows());
        for (int i = 0; i < m_data->rows(); ++i) {
            auto &row = m_data->row(i);
            m_dis(i) = sqrt(row.dot(row));
        }
    }

    int KNNSolver::predict(const Eigen::VectorXd &predict_data) {
        std::map<double, int> sims;
        double self_dis = sqrt(predict_data.dot(predict_data));
        // 重要：跳过空数据防止出错
        if (self_dis < SKIP_VALUE) {
            return -1;
        }
        for (int i = 0; i < m_data->rows(); ++i) {
            if (m_dis(i) < SKIP_VALUE) continue;
            auto &row = m_data->row(i);
            double sim = row.dot(predict_data) / self_dis / m_dis(i);
            if (sim > 0) {
                sims.emplace(sim, i);
            }
        }
        std::unordered_map<int, int> count;
        auto it = sims.rbegin();
        for (int i = 0; i < m_k; ++i) {
            int lb = m_label->operator()(it->second);
            auto p = count.find(lb);
            if (p == count.end()) {
                count.emplace(lb, 1);
            } else {
                ++p->second;
            }
            ++it;
        }
        int lb, mx = 0;
        for (auto p: count) {
            if (p.second > mx) {
                mx = p.second;
                lb = p.first;
            }
        }
        return lb;
    }
} // FansionML
